import json
import time
import logging
import threading
from celery import shared_task
from django.core import serializers
from django.conf import settings
from .models import KnowledgeGraph, UnstructuredFile, KnowledgeGraphStatus, GraphDataSourceTable
from graphNeo4j.lib.Neo4jServices import Neo4jServices
from .services.patent_source_service import PatentSourceService
from dify.api import upload_file as dify_upload_file, run_workflow as dify_run_workflow, parse_workflow_response
from .serializers import (
    GraphDataSourceSerializer,
    GraphDataSourceTableSerializer,
)
from .db_connect import DatabaseService
from functools import wraps
from concurrent.futures import ThreadPoolExecutor, as_completed


logger = logging.getLogger('celery_tasks')

@shared_task(queue='graph')
def process_unstructured_file(graph_id):
    """处理非结构化文件"""
    total_steps = 3
    start_time = time.time()  # 记录开始时间
    try:
        logger.info(f"|-------------------图谱-{graph_id} 非结构化构建开始 -----------------------|")
        logger.info(f"Step 1/{total_steps} 获取图谱非结构化文件")
        graph = KnowledgeGraph.objects.get(id=graph_id)
        graph.status = KnowledgeGraphStatus.BUILDING.value
        graph.save()

        # 文件提取
        files = UnstructuredFile.objects.filter(graph=graph)
        if len(files) == 0:
            raise Exception(f"{graph_id} - 图谱非结构化文件不存在")

        file_list = []
        for file_item in files:
            # 文件上传
            file_id = dify_upload_file('非结构化处理-v2', file_item.file_path, file_item.file_type)
            if file_id:
                file_list.append({
                    'transfer_method': "local_file",
                    'upload_file_id': f"{file_id}",
                    'type': "document"
                })
            else:
                raise Exception(f"{file_item.file_path} - 文件上传失败")
            # break

        # diy解析文件
        logger.info(f"Step 2/{total_steps} Dify 专利文件解析")
        result = dify_run_workflow('非结构化处理-v2', {
            "file": file_list
        })
        # print(result)
        # 解析返回结构
        reponse_str = parse_workflow_response(result)
        if reponse_str == 'False':
            raise Exception(f"构建失败，未从文件提取到图谱数据， file_path:{file_item.file_path}")
        print(reponse_str)
        parse_result = json.loads(reponse_str)
        knowledge_graph_data = parse_result['knowledge_graph']

        print(knowledge_graph_data)

        total_node = 0
        if "entities" in knowledge_graph_data:
            total_node += len(knowledge_graph_data['entities'])

        graph.build_total = total_node
        graph.save()

        # TODO: 实现图谱构建逻辑
        # 检查是否存在关联的数据源
        logger.info(f"Step 3/{total_steps} 提取需要构建和数据表")
        neo4j_config = settings.NEO4J_CONFIG

        # 添加知识节点
        neo4j_services = Neo4jServices(neo4j_config['uri'], neo4j_config['user'], neo4j_config['password'])

        # 添加顶级知识节点
        label1 = f'`{graph.name}`' if '-' in graph.name else graph.name
        # result = neo4j_services.create_node(label1, {'name': graph.name, 'graph_id': graph_id})

        # 添加知识子节点
        nodes = {}
        processed_total = 0
        for entity in knowledge_graph_data['entities']:
            attributes= {
                'name': entity['name'],
                'type': entity['type'],
                'graph_id': graph_id
            }
            for attribute in entity.get('attributes', []):
                attributes[attribute['name']] = attribute['value']


            neo4j_services.create_node_labels([label1, entity['type']], attributes)
            neo4j_services.set_node_labels_props([label1, entity['type']], entity['name'], attributes)

            processed_total += 1
            if total_node <= 1000 and processed_total % 50 == 0:
                update_process(graph_id, processed_total)
            elif total_node > 1000 and processed_total % 100 == 0:
                update_process(graph_id, processed_total)
            elif processed_total > 10000 and processed_total % 1000 == 0:
                update_process(graph_id, processed_total)
            else:
                update_process(graph_id, processed_total)

        print(knowledge_graph_data['relationships'])
        # 添加关系数据
        for relationship in knowledge_graph_data['relationships']:
            print(f"from-{relationship['from']} to-{relationship['to']}")
            neo4j_services.create_relationship_between_nodes(relationship['from'], relationship['to'], label1, label1, relationship['relationship'])

        graph.status = KnowledgeGraphStatus.COMPLETED.value
        graph.already_build = processed_total
        graph.save()

        end_time = time.time()  # 记录结束时间
        elapsed_time = end_time - start_time  # 计算经过的时间
        logger.info(f"|-------------------图谱-{graph_id} 非结构化构建完成 用时：{elapsed_time:.4f} seconds-----------------------|")
    except Exception as e:
        graph.status = KnowledgeGraphStatus.FAILED.value
        graph.save()
        logger.error(f"构建结构化图谱失败: {e}")


@shared_task(queue='graph')
def build_structured_graph(graph_id):
    """构建结构化图谱"""
    total_steps = 3
    start_time = time.time()  # 记录开始时间
    try:
        logger.info(f"|-------------------图谱-{graph_id} 构建开始 -----------------------|")
        logger.info(f"Step 1/{total_steps} 获取图谱数据")
        graph = KnowledgeGraph.objects.get(id=graph_id)
        graph.status = KnowledgeGraphStatus.BUILDING.value
        graph.save()

        already_build = graph.already_build

        # TODO: 实现图谱构建逻辑
        # 检查是否存在关联的数据源
        logger.info(f"Step 2/{total_steps} 提取需要构建和数据表")
        if hasattr(graph, 'data_source'):
            data_source = getattr(graph, 'data_source', None)
            data_source = GraphDataSourceSerializer(data_source).data
        else:
            raise Exception(f"{graph_id} - 图谱数据源不存在")

        if hasattr(graph, 'data_source_table'):
            data_source_tables = getattr(graph, 'data_source_table', None)
            source_tables = GraphDataSourceTableSerializer(data_source_tables, many=True).data
            logger.info(f"Step 3/{total_steps} 逐表数据处理")

            for row in source_tables:
                table_name = row['table_name']
                logger.info(f"  处理数据表: {table_name}")
                # processed_total = loop_write_data(graph.name, graph_id, table_name, data_source, already_build)
                # 多线程处理数据
                processed_total = loop_write_data_thread(graph.name, graph_id, table_name, data_source, already_build)
                # 阶段性更新进度
                update_process(graph_id, processed_total + already_build)

        else:
            raise Exception(f"{graph_id} - 图谱数据表不存在")

        graph.status = KnowledgeGraphStatus.COMPLETED.value
        graph.already_build = processed_total
        graph.save()

        end_time = time.time()  # 记录结束时间
        elapsed_time = end_time - start_time  # 计算经过的时间
        logger.info(
            f"|-------------------图谱-{graph_id} 构建完成 用时：{elapsed_time:.4f} seconds-----------------------|")
    except Exception as e:
        graph.status = KnowledgeGraphStatus.FAILED.value
        graph.save()
        logger.error(f"构建结构化图谱失败: {e}")


# 构建结构化图谱
def loop_write_data(graph_name, graph_id, table_name, data_source, processed_total):
    # 数据库连接
    try:
        db_service = DatabaseService(data_source['db_type'])
        db_service.connection(data_source)

        neo4j_config = settings.NEO4J_CONFIG

        page = 1
        size = 500
        total = db_service.get_table_rows(table_name)
        while (1):
            # if page == 3 or processed_total >= 2000:
            #     break
            data = db_service.get_table_data(table_name, page, size)
            real_size = 0
            for item in data:
                row = PatentSourceService().get_patent_standard_field_data(table_name, item)
                neo4j_services = Neo4jServices(neo4j_config['uri'], neo4j_config['user'], neo4j_config['password'])
                label1 = f'`{graph_name}`' if '-' in graph_name else graph_name
                result = neo4j_services.add_knowledge_to_graph(label1, graph_id, row)
                real_size += 1
                if not result:
                    break

            processed_total += real_size

            # 每1000条更新一次进度
            if processed_total % 100 == 0:
                update_process(graph_id, processed_total)

            logger.info(f"  已处理 {processed_total}/{total} 条")
            if not data:
                break
            page += 1
        return processed_total
    except Exception as e:
        logger.error(f"数据循环处理失败: {e}")
        return

lock = threading.Lock()  # 创建锁对象
# 构建结构化图谱-多线程版
def loop_write_data_thread(graph_name, graph_id, table_name, data_source, processed_total):
    # 数据库连接
    try:
        db_conn = DatabaseService(data_source['db_type'])
        db_conn.connection(data_source)
        page = 1
        size = 500
        with ThreadPoolExecutor(max_workers=10) as executor:  # 根据需要调整线程数
            futures = []
            while True:
                future = executor.submit(db_process_item, data_source, table_name, page, size)
                futures.append(future)
                if db_conn.get_table_rows(table_name) <= page * size:
                    break
                page += 1

            for future in as_completed(futures):
                # 获取任务结果
                data = future.result()
                if data:
                    for item in data:
                        # 并发提交Neo4j写入操作
                        executor.submit(process_item, table_name, graph_name, graph_id, item)
                        with lock:  # 确保对 processed_total 的访问是线程安全的
                            processed_total += 1  # 在这里累加已处理的数量
                            if processed_total % 1000 == 0:
                                update_process(graph_id, processed_total)

        logger.info(f"  已处理 {processed_total}/{page} 条")
        return processed_total
    except Exception as e:
        logger.error(f"数据循环处理失败: {e}")
        return

# 获取数据库数据
def db_process_item(data_source, table_name, page, size):
    try:
        db_service = DatabaseService(data_source['db_type'])
        db_service.connection(data_source)
        data = db_service.get_table_data(table_name, page, size)
        return data
    except Exception as e:
        logger.error(f"获取数据库数据失败: {e}")
        return

def process_item(table_name, graph_name, graph_id, item):
    neo4j_config = settings.NEO4J_CONFIG
    row = PatentSourceService().get_patent_standard_field_data(table_name, item)
    label1 = f'`{graph_name}`' if '-' in graph_name else graph_name
    neo4j_services = Neo4jServices(neo4j_config['uri'], neo4j_config['user'], neo4j_config['password'])
    result = neo4j_services.add_knowledge_to_graph(label1, graph_id, row)
    return result

def update_process(graph_id, processed_total):
    try:
        graph = KnowledgeGraph.objects.get(id=graph_id)
        graph.already_build = processed_total
        graph.save()
    except Exception as e:
        logger.error(f"更新进度失败: {e}")
        return
